from picamera2 import Picamera2
import cv2
from pynput import keyboard
import time
import RPi.GPIO as GPIO

# 硬件引脚定义
MOTOR_PINS = [11, 12, 13, 15]  # IN1, IN2, IN3, IN4
SERVO_PINS = {'lr': 7, 'ud': 18}  # 左右转向, 上下运动

# 全局变量
current_angles = {'lr': 90, 'ud': 90}  # 当前舵机角度

def setup_gpio():
    """初始化GPIO和PWM"""
    GPIO.setmode(GPIO.BOARD)
    GPIO.setwarnings(False)
    
    # 电机引脚初始化
    for pin in MOTOR_PINS:
        GPIO.setup(pin, GPIO.OUT, initial=GPIO.LOW)
    
    # 舵机PWM初始化
    GPIO.setup(SERVO_PINS['lr'], GPIO.OUT)
    GPIO.setup(SERVO_PINS['ud'], GPIO.OUT)
    pwms = {
        'lr': GPIO.PWM(SERVO_PINS['lr'], 50),
        'ud': GPIO.PWM(SERVO_PINS['ud'], 50)
    }
    pwms['lr'].start(0)
    pwms['ud'].start(0)
    return pwms

def set_servo_angle(pwm, angle, servo_type):
    """设置舵机角度（带范围限制）"""
    angle = max(0, min(180, angle))
    duty = angle / 18 + 2
    pwm.ChangeDutyCycle(duty)
    time.sleep(0.15)  # 给舵机转动时间
    pwm.ChangeDutyCycle(0)  # 停止信号防止抖动
    current_angles[servo_type] = angle
    return angle

def main_control(pwms):
    """主控制程序"""
    # 初始化摄像头
    picam2 = Picamera2()
    config = picam2.create_preview_configuration(main={"size": (640, 480), "format": "BGR888"})
    picam2.configure(config)
    picam2.start()

    def motor_control(in1, in2, in3, in4):
        GPIO.output(MOTOR_PINS[0], in1)
        GPIO.output(MOTOR_PINS[1], in2)
        GPIO.output(MOTOR_PINS[2], in3)
        GPIO.output(MOTOR_PINS[3], in4)

    def on_press(key):
        try:
            # 运动控制
            if key.char == 'w': motor_control(1,0,1,0)
            elif key.char == 's': motor_control(0,1,0,1)
            elif key.char == 'a': motor_control(0,1,1,0)
            elif key.char == 'd': motor_control(1,0,0,1)
            
            # 舵机控制（±15°）
            elif key.char == 'q':  # 左转
                set_servo_angle(pwms['lr'], current_angles['lr'] - 15, 'lr')
            elif key.char == 'e':  # 右转
                set_servo_angle(pwms['lr'], current_angles['lr'] + 15, 'lr')
            elif key.char == 'n':  # 上抬
                set_servo_angle(pwms['ud'], current_angles['ud'] - 15, 'ud')
            elif key.char == 'm':  # 下压
                set_servo_angle(pwms['ud'], current_angles['ud'] + 15, 'ud')
            elif key.char == 'h':  # 回正
                set_servo_angle(pwms['lr'], 90, 'lr')
                set_servo_angle(pwms['ud'], 90, 'ud')
                
        except AttributeError:
            pass

    def on_release(key):
        # 松开移动键时停止电机
        if hasattr(key, 'char') and key.char in ['w','a','s','d']:
            motor_control(0,0,0,0)
        # ESC键退出
        if key == keyboard.Key.esc:
            return False

    # 启动控制
    listener = keyboard.Listener(on_press=on_press, on_release=on_release)
    listener.start()

    try:
        print("""\n【控制指令】
    WASD: 移动控制（松开停止）
    Q/E: 左右舵机±15° 
    N/M: 上下舵机±15°
    H: 舵机回正(90°)
    ESC: 退出程序""")
        
        while listener.is_alive():
            frame = picam2.capture_array()
            cv2.imshow("Camera", frame)
            if cv2.waitKey(1) == 27:  # ESC键ASCII码
                break
    finally:
        # 清理资源
        motor_control(0,0,0,0)
        for pwm in pwms.values():
            pwm.ChangeDutyCycle(0)
            pwm.stop()
        picam2.stop()
        listener.stop()
        cv2.destroyAllWindows()
        GPIO.cleanup()
        print("程序已安全退出")

if __name__ == "__main__":
    pwms = setup_gpio()
    try:
        # 启动时自动归中
        print("初始化舵机到中间位置...")
        set_servo_angle(pwms['lr'], 90, 'lr')
        set_servo_angle(pwms['ud'], 90, 'ud')
        
        # 启动主控制
        main_control(pwms)
    except Exception as e:
        print(f"错误: {str(e)}")
        GPIO.cleanup()